class LightningModel(pl.LightningModule):
def __init__(self, model, optimizer_name, optimizer_hparams):
super().__init__()
self.model = model
self.save_hyperparameters(ignore=['model'])
def _random_bbox(self, n):
values = list(range(-5, 6))
random_vectors_list = [random.sample(values, 4) for _ in range(n)]
return [torch.tensor([10, 30, 170, 190], dtype=torch.float, device=device) + torch.tensor(rv, dtype=torch.float, device=device) for rv in random_vectors_list]
def _data_step(self, batch):
x, (y_lbl, y_kpt) = batch
y_lbl += 1
y_kpt = y_kpt.view(-1, 5, 2)
y_kpt = torch.cat((y_kpt, torch.ones((y_kpt.shape[0], 5, 1), device=device)), dim=2)
bboxs = self._random_bbox(len(x))
y = [{'boxes': bbox[None], 'labels': lbl[None], 'keypoints': kpt[None]} for bbox, lbl, kpt in zip(bboxs, y_lbl, y_kpt)]
return x, y
def _train_model(self, x, y):
with torch.no_grad():
self.model.train()
return self.model(x, y)
def forward(self, x, y=None):
return self.model(x, y)
def training_step(self, batch, batch_idx):
x, y = self._data_step(batch)
losses = self(x, y)
loss = sum([loss for loss in losses.values()])
self.log_dict({'train_loss': loss}, on_epoch=True, on_step=False, prog_bar=True)
return loss
def validation_step(self, batch, batch_idx):
x, y = self._data_step(batch)
losses = self._train_model(x, y)
loss = sum([loss for loss in losses.values()])
self.log_dict({'valid_loss': loss}, on_epoch=True, on_step=False, prog_bar=True)
losses = self(x, y)
def test_step(self, batch, batch_idx):
x, y = self._data_step(batch)
losses = self._train_model(x, y)
loss = sum([loss for loss in losses.values()])
self.log_dict({'test_loss': loss}, on_epoch=True, on_step=False, prog_bar=True)
losses = self(x, y)
def configure_optimizers(self):
# We will support Adam or SGD as optimizers.
if self.hparams.optimizer_name == "Adam":
optimizer = torch.optim.AdamW(self.parameters(), **self.hparams.optimizer_hparams)
elif self.hparams.optimizer_name == "SGD":
optimizer = torch.optim.SGD(self.parameters(), **self.hparams.optimizer_hparams)
elif self.hparams.optimizer_name == 'ASGD':
optimizer = torch.optim.ASGD(self.parameters(), **self.hparams.optimizer_hparams)
else:
assert False, f'Unknown optimizer: "{self.hparams.optimizer_name}"'
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.3)
return {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "valid_loss"}
